import torch
from dataset import cancerClassDataset
import class8dataset
from math import sqrt
import torchvision

def calculate_means_STDs(dataset:cancerClassDataset):
    num_imgs=len(dataset)
    means_channel=torch.zeros(size=(3,num_imgs,))
    for i,(raw_img,_) in enumerate(dataset):
        if i%50==0:
            print('进度:',i)
        assert(len(raw_img)==3)
        means_channel[0][i]=torch.mean(raw_img[0])
        means_channel[1][i]=torch.mean(raw_img[1])
        means_channel[2][i]=torch.mean(raw_img[2])
    
    mean1,mean2,mean3=torch.mean(means_channel[0]),\
        torch.mean(means_channel[1]),torch.mean(means_channel[2])
    

    square_means=torch.zeros(size=(3,num_imgs,))
    for i,(raw_img,_) in enumerate(dataset):
        if i%50==0:
            print('进度:',i)
        assert(len(raw_img)==3)
        square_means[0][i]=torch.mean(raw_img[0]**2)
        square_means[1][i]=torch.mean(raw_img[1]**2)
        square_means[2][i]=torch.mean(raw_img[2]**2)
    std1=torch.mean(square_means[0])-mean1**2
    std2=torch.mean(square_means[1])-mean2**2
    std3=torch.mean(square_means[2])-mean3**2

    print('平均值:',mean1.item(),mean2.item(),mean3.item())
    print('方差:',std1,std2,std3)
    print('标准差:',sqrt(std1),sqrt(std2),sqrt(std3))

if __name__ == '__main__':
    allData=class8dataset.getRandomDataPaths('E:/JohnsonProj/Kather_texture_2016_image_tiles_5000')
    set=class8dataset.cancer5000Dataset(allData)
    calculate_means_STDs(set)

if __name__ == '__main1__':
    means=(0.5843984484672546, 0.4723328649997711, 0.6498987674713135,)
    stds=(0.267836341731424, 0.3274180198302301, 0.25589914959136034,)
    allData=class8dataset.getRandomDataPaths('E:/JohnsonProj/Kather_texture_2016_image_tiles_5000')
    set=class8dataset.cancer5000Dataset(allData,\
        raw_transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(means,stds),
        ])
    )
    calculate_means_STDs(set)


if __name__ == '__main1__':
    set=cancerClassDataset('E:/JohnsonProj/smallTraining',)

    calculate_means_STDs(set)

if __name__ == '__main1__':
    means=[0.7084528207778931, 0.5368406772613525, 0.7413279414176941]
    stds=[0.1579459655207818, 0.21840702144406327, 0.16517678558091184]
    set=cancerClassDataset('E:/JohnsonProj/smallTesting',raw_transform=
    torchvision.transforms.Compose(
        [
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(means,stds),
        ]
    ))

    calculate_means_STDs(set)